#!/usr/bin/python
# Audio converter script
# Copyright 2007-2011 Mkhanyisi Madlavana

import time
import optparse
import subprocess as sub

import mutagen
import mutagen.id3 as id3
import mutagen.mp3 as mp3
import mutagen.mp4 as mp4
import mutagen.asf as asf
import mutagen.flac as flac
import mutagen.oggvorbis as ogg

def get_ext(f):
    """ 
    returns (filename_without_ext, ext)
    eg: 'mysong.mp3' will be returned as 
    ('mysong','.mp3')
    """
    idx = f.rfind('.')
    if idx < 0:
        return (f, None)
    else:
        return (f[:idx], f[idx:])

def set_metadata(files, mytags):
    """
    sets metadata tags as specified by the 'mytags' dictionary
    keys/tags with empty values in 'mytags' are ignored.
    NB: the 'files' variable must be a list.
    """
    for f in files:
        dummy, ext = get_ext(f)
        if ext == '.mp3':
            metadata = mp3.Open(f)
            tags = {}
            if mytags['artist'] is not None:
                tags['TPE1'] = id3.TPE1(encoding=3, text=[u'%s' % mytags['artist']])
            if mytags['title'] is not None:
                tags['TIT2'] = id3.TIT2(encoding=3, text=[u'%s' % mytags['title']])
            if mytags['album'] is not None:
                tags['TALB'] = id3.TALB(encoding=3, text=[u'%s' % mytags['album']])
            if mytags['year'] is not None:
                tags['TDRC'] = id3.TDRC(encoding=3, text=[u'%s' % mytags['year']])
            if mytags['genre'] is not None:
                tags['TCON'] = id3.TCON(encoding=3, text=[u'%s' % mytags['genre']])
            if mytags['tracknumber'] is not None:
                tags['TRCK'] = id3.TRCK(encoding=3, text=[u'%s' % mytags['tracknumber']])

        elif ext == '.wma':
            metadata = asf.Open(f)
            tags = {}
            if mytags['title'] is not None:
                tags["Title"] = mytags['title']
            if mytags['artist'] is not None:
                tags["Author"] = mytags['artist']
            if mytags['album'] is not None:
                tags["WM/AlbumTitle"] = mytags['album']
            if mytags['year'] is not None:
                tags["WM/Year"] = mytags['year']
            if mytags['genre'] is not None:
                tags["WM/Genre"] = mytags['genre']
            if mytags['tracknumber'] is not None:
                tags["WM/TrackNumber"] = mytags['tracknumber']

        elif ext in ('.ogg', '.flac'):
            if ext == '.flac':
                metadata = flac.Open(f)
            else:
                metadata = ogg.Open(f)
            tags = {}
            for key, value in mytags.items():
                if value is not None:
                    if key == "year":
                        key = "date"
                    tags[key] = value

        elif ext in ('.m4a', '.aac', '.mp4'):
            metadata = mp4.Open(f)
            tags = {}
            if mytags['title'] is not None:
                tags['\xa9nam'] = [u'%s' % mytags['title']]
            if mytags['artist'] is not None:
                tags['\xa9ART'] = [u'%s' % mytags['artist']]
            if mytags['album'] is not None:
                tags['\xa9alb'] = [u'%s' % mytags['album']]
            if mytags['genre'] is not None:
                tags['\xa9gen'] = [u'%s' % mytags['genre']]
            if mytags['year'] is not None:
                tags['\xa9day'] = [u'%s' % mytags['year']]
            if mytags['tracknumber'] is not None:
                # tracknumber for track 8, disk 2 is 8/2 (ie tracknum/disknum)
                try:
                    t = map(int,mytags['tracknumber'].split('/'))
                    tags['trkn'] = [(t[0], t[1])]
                except IndexError:
                    try:
                        tags['trkn'] = [(t[0],0)]
                    except:
                        tags['trkn'] = [(mytags['tracknumber'],0)]
                except:
                    if tags.has_key('trkn'):
                        tags.pop('trkn')
        else:
            continue

        for key, value in tags.items():
            if value is not None:
                metadata[key] = value
        metadata.save()

def get_metadata(files):
    """
    gets, and prints out if required, the song information (aka metadata)
    the 'files' variable must be a list.
    """
    for f in files:
        artist, album, title = None, None, None
        genre, year, tracknumber = None, None, None
        ext = get_ext(f.lower())[1]
        metadata = mutagen.File(f, easy=True)

        if ext in ['.ogg', '.flac', '.mp3', '.mp4', '.aac', '.m4a']:
            if metadata.has_key('artist'):
                artist = metadata['artist'][0]
            if metadata.has_key('album'):
                album = metadata['album'][0]
            if metadata.has_key('title'):
                title = metadata['title'][0]
            if metadata.has_key('genre'):
                genre = metadata['genre'][0]
            if metadata.has_key('date'):
                year = metadata['date'][0]
            if metadata.has_key('tracknumber'):
                tracknumber = metadata['tracknumber'][0]
        elif ext == '.wma':
            if metadata.has_key('Author'):
                artist = metadata['Author'][0]
            if metadata.has_key('WM/AlbumTitle'):
                album = metadata['WM/AlbumTitle'][0]
            if metadata.has_key('Title'):
                title = metadata['Title'][0]
            if metadata.has_key('WM/Genre'):
                genre = metadata['WM/Genre'][0]
            if metadata.has_key('TRCK'):
                year = metadata['WM/Year'][0]
            if metadata.has_key('WM/TrackNumber'):
                tracknumber = metadata['WM/TrackNumber'][0]

        metadata_tags = {
                "title"         : title, 
                "artist"        : artist, 
                "album"         : album, 
                "tracknumber"   : tracknumber, 
                "genre"         : genre, 
                "year"          : year
                }

        nonempty_count = 0
        for key, val in metadata_tags.items():
            if val is not None:
                nonempty_count += 1
                try:
                    metadata_tags[key] = unicode(val,'utf-8')
                except:
                    continue
        return (metadata_tags, nonempty_count)

def log_time():
    """
    returns current time in HOUR:MINUTE:SECOND format
    """
    return time.strftime("%H:%M:%S", time.localtime())

def check_deps(check=False):
    """
    checks whether all dependencies for this script are installed or not.
    the results are printed out on the screen if check=True or are redirected 
    to a log file if check=False
    """

    # deps = {dep1: [audio formats needing it], dep2: [audio formats needing it]}
    deps = {
            'mutagen-inspect'   : [],
            'ffmpeg'            : ['mp3', 'ogg', 'wma', 'm4a', 'flac', 'wav'],
            'lame'              : ['mp3'],
            'flac'              : ['flac'],
            'faac'              : ['m4a'],
            'oggenc'            : ['ogg'],
            }

    for d in deps.keys():
        pkg = sub.Popen(["which", d], stdout=sub.PIPE).communicate()[0].strip() 
        if check:
            print d + '...',
            if not pkg:
                print red + "not installed" + nc + "."
            else:
                print green + "installed" + nc + "."
        else:
            log_file.write("[%s]" % log_time())
            if pkg:
                log_file.write("%s\n" % pkg)
            else:
                msg = "_______ %s not installed _______" % d
                print red + msg + nc + "."
                log_file.write(msg + "\n")
                for fmt in deps[d]:
                    quality_presets.pop(fmt)
    if check:
        raise SystemExit(0)

def convert_to_mp3(filename, ext, stderr_out, preset=None):
    """
    converts input track to MPEG-1 Layer III format (aka MP3)
    """
    cmd1 = ["ffmpeg", "-y", "-i", filename + ext, "-f", "wav", "/dev/stdout"]
    cmd2 = ("lame " + preset).split()
    cmd2.extend(["-", "./" + filename + ".mp3"])
    return convert(stderr_out, cmdline1=cmd1, cmdline2=cmd2)

def convert_to_m4a(filename, ext, stderr_out, preset=None):
    """
    converts input track into MPEG-4 audio format (aka AAC/M4A/MP4)
    """
    cmd1 = ["ffmpeg", "-y", "-i", filename + ext, "-f", "wav", "/dev/stdout"]
    cmd2 = ("faac " + preset).split()
    cmd2.extend(["-o", "./" + filename + ".m4a", "/dev/stdin"])
    return convert(stderr_out, cmdline1=cmd1, cmdline2=cmd2)
   
def convert_to_wma(filename, ext, stderr_out, preset=None):
    """
    converts an input track into the Microsoft ASF format (aka WMA)
    """
    cmd1 = ["ffmpeg", "-y", "-i", filename + ext]
    cmd2 =  preset.split()
    cmd1.extend(cmd2)
    cmd1.append("./" + filename + ".wma")
    return convert(stderr_out, cmdline1=cmd1)

def convert_to_ogg(filename, ext, stderr_out, preset=None):
    """
    converts input track into Ogg Vorbis format
    """
    cmd1 = ["ffmpeg", "-y", "-i", filename + ext, "-f", "wav", "/dev/stdout"]
    cmd2 = ("oggenc " + preset).split()
    cmd2.extend(["-o", "./" + filename + ".ogg", "/dev/stdin"])
    return convert(stderr_out, cmdline1=cmd1, cmdline2=cmd2)
   
def convert_to_wav(filename, ext, stderr_out, preset=None):
    """
    decodes a track into WAVE format. This format has no metadata support.
    """
    cmd = ["ffmpeg", "-y", "-i", filename + ext, "-f", "wav", "./" + filename + ".wav"]
    return convert(stderr_out, cmdline1=cmd)

def convert_to_flac(filename, ext, stderr_out, preset=None):
    """
    converts input track into Free Lossless Audio Codec format (aka FLAC).
    this format is lossless and supports metadata tags.
    """
    status = convert_to_wav(filename, ext, stderr_out)
    if status:
        return 1

    cmd = ("flac " + preset).split()
    cmd.extend(["-f", "-o", "./" + filename + ".flac", filename + ".wav"])
    return convert(stderr_out, cmdline1=cmd)

def convert(stderr_out,cmdline1=None, cmdline2=None):
    if cmdline1 is not None and cmdline2 is not None:
        p1 = sub.Popen(cmdline1, stdout=sub.PIPE, stderr=stderr_out)
        p2 = sub.Popen(cmdline2, stdin=p1.stdout, stdout=sub.PIPE, stderr=stderr_out)
        output = p2.communicate()
        return p2.returncode
    elif cmdline1 is not None and cmdline2 is None:
        p1 = sub.Popen(cmdline1, stdout=sub.PIPE, stderr=stderr_out)
        output = p1.communicate()
        return p1.returncode
    else:
        raise SystemExit("Error: unexpected arguments on the convert() function")

# audio quality presets.
quality_presets = {
        "mp3": 
        {
            "insane"    : "--cbr -b 320",
            "extreme"   : "-V 0",
            "high"      : "-V 2",
            "normal"    : "-V 4",
            "low"       : "--preset 64",
            "tiny"      : "--cbr -b 32 -m m -s 32",
            },
        "ogg": 
        {
            "insane"    : "-q 10",
            "extreme"   : "-q 8",
            "high"      : "-q 6",
            "normal"    : "",
            "low"       : "-q 1",
            "tiny"      : "-q -1",
            },
        "wma":
        {
            "insane"    : "-ab 320000 -ac 2 -acodec wmav2",
            "extreme"   : "-ab 256000 -ac 2 -acodec wmav2",
            "high"      : "-ab 192000 -ac 2 -acodec wmav2",
            "normal"    : "-ab 128000 -ac 2 -acodec wmav2",
            "low"       : "-ab 64000 -ac 2 -acodec wmav2",
            "tiny"      : "-ab 32000 -ac 2 -acodec wmav2",
            },
        "m4a":
        {
            "insane"    : "-q 500 -w",
            "extreme"   : "-q 350 -w",
            "high"      : "-q 200 -w",
            "normal"    : "-w",
            "low"       : "-q 50 -w",
            "tiny"      : "-q 10 -w",
            },
        "flac":
        {
            "insane"    : "--best",
            "extreme"   : "-7",
            "high"      : "-6",
            "normal"    : "-5",
            "low"       : "-3",
            "tiny"      : "--fast",
            },
        "wav": 
        {
            "insane"    : None, 
            "extreme"   : None, 
            "high"      : None,
            "normal"    : None, 
            "low"       : None, 
            "tiny"      : None,
            }
}

converters = {
        "mp3"   : convert_to_mp3,
        "ogg"   : convert_to_ogg,
        "m4a"   : convert_to_m4a,
        "wma"   : convert_to_wma,
        "wav"   : convert_to_wav,
        "flac"  : convert_to_flac,
        }

if __name__ == "__main__":
    parser = optparse.OptionParser(usage="%prog [options] [files]", version="4.0.1")
    parser.add_option('-f', '--format', type=str, default='mp3', help='audio format to convert to')
    parser.add_option('-q', '--quality', type=str, default='normal', help='audio quality preset')
    parser.add_option('-c', '--check', dest='check', action='store_true', help='check dependencies')
    parser.add_option('-r', '--remove', dest='remove', action='store_true', help='remove original file after converting successfully')
    parser.add_option('-d', '--decode', dest='decode', action='store_true', help='decode file .wav format')
    parser.add_option('-w', '--over', dest='overwrite', action='store_true', help='overwrite destination file if it exists already')
    parser.add_option('-u', '--unlock', dest='unlock', action='store_true', help='unlock a locked file and convert')
    parser.add_option('--directory', dest="walk", default=None, type=str, help='convert all files inside the given directory')
    opt, files = parser.parse_args()

    if len(files) < 1 and opt.walk is None and opt.check is None:
        raise SystemExit("ftransc: no input file")
    
    if opt.walk is not None:
        walker = sub.os.walk(opt.walk)
        dest_dir, dummy, files = walker.next()
        pwd = sub.os.getcwd()
        sub.os.chdir(dest_dir)
        
    # print colors
    red = "\033[1;31m"
    green = "\033[1;32m"
    blue = "\033[1;34m"
    pink = "\033[1;35m"
    nc = "\033[0m"
   
    opt.quality = opt.quality.lower()
    opt.format = opt.format.lower()
    if opt.format in ("mp4","m4a","aac"):
        opt.format = "m4a"
    if opt.decode:
        opt.format = "wav"
    if not quality_presets.has_key(opt.format):
        raise SystemExit("%s%s%s is not a supported format" % (red, opt.format, nc))
    if not quality_presets[opt.format].has_key(opt.quality):
        print "%s%s%s is not a valid quality preset, using %s%s%s." % (red, opt.quality, nc, green, 'normal', nc)
        opt.quality = "normal"
    if not opt.check:
        log_file = open('/tmp/ftransc.log','w')
        null_file = open('/dev/null','w')
    check_deps(opt.check)
    if not quality_presets.has_key(opt.format):
        raise SystemExit("%s___ %s ___%s format not supported. required dependency not installed." % (red, opt.format, nc))

    # go through each file: 1. extract tags, 2. convert, and 3. insert tags
    count = 0
    for in_file in files:
        count += 1
        total = len(files)
        out_file = get_ext(in_file)[0] + "." + opt.format
        if out_file == in_file:
            print "[%s%d/%d%s]: %s%s%s: same input file & output file" % (pink, count, total, nc, red, in_file, nc)
            continue
        if not sub.os.path.exists(in_file):
            print "[%s%d/%d%s]: %s%s%s does not exist" % (pink, count, total, nc, red, in_file, nc)
            continue
        if sub.os.path.isfile(out_file) and not opt.overwrite:
            print "[%s%d/%d%s]: skipping %s%s%s: overwrite disabled" % (pink, count, total, nc, red, in_file, nc)
            continue
        if sub.os.path.isdir(in_file) and opt.walk is None:
            print "[%s%d/%d%s]: skipping directory %s%s%s, use '--directory'" % (pink, count, total, nc, red, in_file, nc)
            continue
        swp_file = ".%s.swp" % in_file
        if sub.os.path.isfile(swp_file):
            if opt.unlock:
                pass
            else:
                print "[%s%d/%d%s]: '%s' is %slocked%s" % (pink, count, total, nc, in_file, red, nc)
                continue
        else:
            try:
                dummy = open(swp_file ,"w")
                dummy.close()
            except:
                raise SystemExit("You %sdo not have permissions%s to write to this folder" % (red, nc))

        try:
            log_file.write("\n[%s]extracting metadata from '%s' ... " % (log_time(), in_file))
            metadata, tag_count = get_metadata([in_file])
            if tag_count < 1:
                log_file.write("no metadata found\n")
            else:
                log_file.write('OK\n')  

        except IOError:
            print "[%s%d/%d%s]: Unable to read %s%s%s" % (pink, count, total, nc, red, in_file, nc)
            log_file.write('Fail\n')
            dummy = sub.os.remove(swp_file)
            continue
        
        in_filename, in_ext = get_ext(in_file)
        log_file.write("[%s]converting '%s' to %s format ... " % (log_time(), in_file, opt.format.upper()))
        if not converters[opt.format](in_filename, in_ext, null_file, preset=quality_presets[opt.format][opt.quality]):
            print "[%s%d/%d%s]: '%s' => %s format = %sSuccess%s" % (pink, count, total, nc, in_file, opt.format.upper(), green, nc)
            log_file.write('OK\n')
            if opt.remove:
                dummy = sub.os.remove(in_file)
            dummy = sub.os.remove(swp_file)
        else:
            print "[%s%d/%d%s]: '%s' => %s format = %sFail%s" % (pink, count, total, nc, in_file, opt.format.upper(), red, nc)
            log_file.write('Fail\n')
            dummy = sub.os.remove(swp_file)
            continue
        
        if tag_count > 0:
            try:
                log_file.write("[%s]inserting metadata to '%s' ... " % (log_time(), out_file))
                set_metadata([out_file], metadata)
                log_file.write("OK\n")
            except:
                print "%sUnable%s to insert metadata to '%s'" % (red, nc, out_file)
                log_file.write("Fail\n")
                continue
       
    if not opt.check:
        log_file.close()
        null_file.close()
